import asyncio
import os
import time

from agents import Runner, WebSearchTool, ModelSettings
from pydantic import BaseModel
# from rich.jupyter import display

from src.common_tools import getAgent

####################################规划Agent#################################################
PROMPT = (
    "You are a helpful research assistant. Given a query, come up with a set of web searches to perform to best answer \ "
    "the query. Output between 20 and 30 terms to query for."
)

class WebSearchItem(BaseModel):

    query: str
    """The search term to use for the web search."""


    reason: str
    """Your reasoning for why this search is important to the query."""

class WebSearchPlan(BaseModel):

    searches: list[WebSearchItem]
    """A list of web searches to perform to best answer the query."""

planner_agent = getAgent(agent_name="web search planner", instructions=PROMPT, tools=[],handoff=[],provider="OPEN_ROUTER")
planner_agent.output_type = WebSearchPlan

####################################那家Agent#################################################

INSTRUCTIONS = (
    "You are a research assistant. Given a search term, you search the web for that term and "
    "produce a concise summary of the results. The summary must 2-3 paragraphs and less than 300 "
    "words. Capture the main points. Write succinctly, no need to have complete sentences or good "
    "grammar. This will be consumed by someone synthesizing a report, so its vital you capture the "
    "essence and ignore any fluff. Do not include any additional commentary other than the summary "
    "itself."
)

search_agent = getAgent(
    agent_name="Search agent",
    instructions=INSTRUCTIONS,
    tools=[WebSearchTool()],
    model_settings=ModelSettings(tool_choice="required"),
    provider="OPEN_ROUTER"
)

####################################书写Agent#################################################
PROMPT = (
    "You are a senior researcher tasked with writing a cohesive report for a research query. "
    "You will be provided with the original query, and some initial research done by a research "
    "assistant.\n"
    "You should first come up with an outline for the report that describes the structure and "
    "flow of the report. Then, generate the report and return that as your final output.\n"
    "The final output should be in markdown format, and it should be lengthy and detailed. Aim "
    "for 5-10 pages of content, at least 1000 words."
)


class ReportData(BaseModel):
    short_summary: str
    """A short 2-3 sentence summary of the findings."""

    markdown_report: str
    """The final report"""

    follow_up_questions: list[str]
    """Suggested topics to research further"""


writer_agent = getAgent(
    agent_name="WriterAgent",
    instructions=PROMPT,
    provider="OPEN_ROUTER"
)
writer_agent.output_type = ReportData


#####################主类########################
from rich.markdown import Markdown


class ResearchManager:
    def __init__(self):
        print("初始化已完成，欢迎使用。使用前请确认相关模型能够被顺利调用。")

    async def run(self, query: str) -> None:
        # Indicate that the research process is starting
        print("Starting research...")

        # Step 1: Generate search plan using planner_agent
        search_plan = await self._plan_searches(query)
        # Step 2: Perform the searches using search_agent
        search_results = await self._perform_searches(search_plan)
        # Step 3: Write the final report using writer_agent
        report = await self._write_report(query, search_results)

        # Final printed report
        print("\n\n=====REPORT=====\n\n")

        print(report.markdown_report)
        print("\n\n=====FOLLOW UP QUESTIONS=====\n\n")
        follow_up_questions = "\n".join(report.follow_up_questions)
        print(follow_up_questions)

        # 保存为 Markdown 文件
        self.save_report_as_md(query, report.markdown_report)

    async def _plan_searches(self, query: str) -> WebSearchPlan:
        print("Planning searches...")
        result = await Runner.run(
            planner_agent,
            f"Query: {query}",
        )
        return result.final_output_as(WebSearchPlan)

    async def _perform_searches(self, search_plan: WebSearchPlan) -> list[str]:
        print("Starting searching...")
        num_completed = 0
        tasks = [asyncio.create_task(self._search(item)) for item in search_plan.searches]
        results = []
        for task in asyncio.as_completed(tasks):
            result = await task
            if result is not None:
                results.append(result)
            num_completed += 1
            print(f"Searching... {num_completed}/{len(tasks)} completed")
        return results

    async def _search(self, item: WebSearchItem) -> str | None:
        print(f"Search term: {item.query}\nReason for searching: {item.reason}")
        try:
            result = await Runner.run(
                search_agent,
                input=f"Search term: {item.query}\nReason for searching: {item.reason}"
            )
            return str(result.final_output)
        except Exception:
            return None

    async def _write_report(self, query: str, search_results: list[str]) -> ReportData:
        print("Thinking about report...")
        print(f"Original query: {query}\nSummarized search results: {search_results}")

        # Use run instead of run_streamed to avoid async context issues
        result = await Runner.run(
            writer_agent,
            input=f"Original query: {query}\nSummarized search results: {search_results}",
        )

        update_messages = [
            "Thinking about report...",
            "Planning report structure...",
            "Writing outline...",
            "Creating sections...",
            "Cleaning up formatting...",
            "Finalizing report...",
            "Finishing report...",
        ]

        last_update = time.time()
        next_message = 0
        for _ in result.new_items:
            if time.time() - last_update > 5 and next_message < len(update_messages):
                print(update_messages[next_message])
                next_message += 1
                last_update = time.time()

        return result.final_output_as(ReportData)

    def save_report_as_md(self, query: str, markdown_content: str) -> None:
        """
        保存生成的报告为 Markdown 文件
        """
        # 创建文件夹（如果不存在）
        folder_name = "research_reports"
        if not os.path.exists(folder_name):
            os.makedirs(folder_name)

        # 使用用户的查询作为文件名
        sanitized_query = query.replace(" ", "_").replace("：", "").replace("?", "")
        file_name = f"{folder_name}/关于{sanitized_query}调研报告.md"

        # 写入 Markdown 文件
        with open(file_name, "w", encoding="utf-8") as file:
            file.write(markdown_content)

        print(f"Report saved as: {file_name}")

if __name__ == "__main__":
    manager = ResearchManager()
    asyncio.run(manager.run("AI在教育中的应用"))

